package com.googlecode.gaal.data.impl;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.NavigableMap;
import java.util.Set;
import java.util.TreeMap;
import java.util.Map.Entry;

import com.googlecode.gaal.preprocess.api.Tokenizer;
import com.googlecode.gaal.preprocess.api.Tokenizer.Document;

public class TreeMapCorpus extends AbstractCorpus<String> {
    private final Map<String, Integer> tokenIntMap = new TreeMap<String, Integer>();
    private final NavigableMap<Integer, Integer> positionDocumentMap = new TreeMap<Integer, Integer>();
    private final String[] intTokenArray;
    private final int[] text;
    private final int alphabetSize;

    public TreeMapCorpus(Tokenizer<String> tokenizer) {
        this(tokenizer, (Set<String>) null, -1);
    }

    public TreeMapCorpus(Tokenizer<String> tokenizer, Set<String> separators) {
        this(tokenizer, separators, -1);
    }

    public TreeMapCorpus(Tokenizer<String> tokenizer, Set<String> separators, int maxDocuments) {
        if (separators == null)
            separators = Collections.emptySet();
        List<Integer> textList = new ArrayList<Integer>();
        boolean isFirst = true;
        for (Document<String> document : tokenizer) {
            if (isFirst)
                isFirst = false;
            else
                textList.add(-1);
            positionDocumentMap.put(textList.size(), document.getId());
            if (maxDocuments != -1 && maxDocuments < document.getId())
                break;
            for (String token : document) {
                Integer tokenIndex;
                if (separators.contains(token)) {
                    tokenIndex = -1;
                } else {
                    tokenIndex = tokenIntMap.get(token);
                    if (tokenIndex == null) {
                        tokenIndex = tokenIntMap.size() + 1;
                        tokenIntMap.put(token, tokenIndex);
                    }
                }
                textList.add(tokenIndex);
            }
        }
        this.text = new int[textList.size()];
        int specialSymbolsCounter = tokenIntMap.size() + 1;
        for (int i = 0; i < textList.size(); i++) {
            int token = textList.get(i);
            if (token == -1) {
                this.text[i] = specialSymbolsCounter++;
            } else {
                this.text[i] = token;
            }
        }
        intTokenArray = new String[tokenIntMap.size() + 1];
        for (Map.Entry<String, Integer> entry : tokenIntMap.entrySet()) {
            intTokenArray[entry.getValue()] = entry.getKey();
        }
        alphabetSize = specialSymbolsCounter - 1;
    }

    /**
     * This constructor is intended to be used when a full control over the
     * alphabet is required, the positions in the supplied alphabet will be
     * used. In case a token in the text is not in the alphabet an exception is
     * thrown
     * 
     * @param corpus
     *            an iterator over the corpus
     * @param alphabet
     *            a list of the types used in the corpus, the zeroth element is
     *            reserved and can not be used
     */
    public TreeMapCorpus(Tokenizer<String> tokenizer, String[] alphabet) {
        for (int i = 1; i < alphabet.length; i++) {
            tokenIntMap.put(alphabet[i], i);
        }
        List<Integer> textList = new ArrayList<Integer>();
        for (Document<String> document : tokenizer) {
            positionDocumentMap.put(textList.size(), document.getId());
            for (String token : document) {
                Integer tokenIndex = tokenIntMap.get(token);
                if (tokenIndex == null) {
                    throw new NullPointerException("symbol '" + token + "' is not in the alphabet");
                }
                textList.add(tokenIndex);
            }
        }
        this.text = new int[textList.size()];
        for (int i = 0; i < textList.size(); i++)
            this.text[i] = textList.get(i);
        intTokenArray = alphabet;
        alphabetSize = intTokenArray.length - 1;
    }

    @Override
    public ArraySequence sequence() {
        return new ArraySequence(text);
    }

    @Override
    public Set<String> alphabet() {
        return tokenIntMap.keySet();
    }

    @Override
    public int alphabetSize() {
        return alphabetSize;
    }

    @Override
    public int toInt(String token) {
        Integer i = tokenIntMap.get(token);
        if (i == null)
            throw new IllegalArgumentException("token not in the corpus");
        return i;
    }

    @Override
    public String toToken(int i) {
        if (i >= intTokenArray.length)
            return Integer.toString(i);
        else
            return intTokenArray[i];
    }

    @Override
    public boolean isSeparator(int symbol) {
        return symbol >= intTokenArray.length;
    }

    @Override
    public int getDocumentId(int index) {
        Entry<Integer, Integer> entry = positionDocumentMap.floorEntry(index);
        if (index < 0 || index >= text.length)
            throw new IndexOutOfBoundsException(String.format("index:%d, text length:%d", index, text.length));
        else
            return entry.getValue();
    }
}
